import torch
from torch import nn
class LSTM_v1(nn.Module):
    def __init__(self):
        super(LSTM_v1, self).__init__()
        INT_HIDDEN_SIZE = 32
        INT_NUM_LAYERS = 1
        CONST_INT_INPUT_SIZE = 5  # [x, y, theta, lstim, rstim]
        CONST_INT_OUTPUT_SIZE = 3  # [x, y, theta]
        CONST_BOOL_BATCH_FIRST = True
        
        self.lstm = nn.LSTM(
            input_size=CONST_INT_INPUT_SIZE,
            hidden_size=INT_HIDDEN_SIZE,
            num_layers=INT_NUM_LAYERS,
            batch_first=CONST_BOOL_BATCH_FIRST
        )
        self.linear = nn.Linear(INT_HIDDEN_SIZE, CONST_INT_OUTPUT_SIZE)

    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        last_output = lstm_out[:, -1, :]
        output = self.linear(last_output)
        return output

def get_model():
    return LSTM_v1()